iT邦幫忙

2023 iThome 鐵人賽

2

簡介

Activation-Aware Weight Quantization (AWQ) 是類似於 GPTQ 的另外一種量化方法,同樣也是透過少量的校準資料集來進行更精準的量化,但是其量化模型的過程所需花費的時間,相較於 GPTQ 而言快了很多。量化後的模型結構,在推論階段時可以有相當快的速度。在社群上也有類似 AutoGPTQ 的 AutoAWQ 這種方便的工具可用,大大簡化了整體量化的流程。今天就來介紹 AWQ 的概念與如何使用 AutoAWQ 進行量化。

原理概念

AWQ 的論文指出,像是 GPTQ 這種透過反向傳播來進行量化的方法,可能會使 LLM 在校準資料集上發生 Overfitting 的情形。因此 AWQ 並不仰賴反向傳播,只透過前向運算來觀察輸出的分佈 (Activation Distribution) 並尋找重要的模型權重,這些權重在論文中被稱為 Salient Weights。作者指出,只要將其中 1% 的重要權重保留為 FP16 的資料型態,就能大幅降低量化帶來的誤差。

但如果為了保護這些分散在不同地方的重要權重,而使得整份權重量化完後是混精度的,則會使計算效率降低。因此作者引入 Activation-Aware Scaling 的技巧來對每組權重做量化,使得重要權重能夠獲得保護的同時,也能夠避免混精度權重以提昇硬體友善度,降低運算瓶頸,在實際推論上可以有相當快的速度。

GEMV vs GEMM

AWQ 分成 GEMV 與 GEMM 兩種矩陣乘法。因為文件沒有特別提到,所以筆者稍微雲了一下,GEMV 應該是指 General Matrix Vector Multiply,是矩陣對向量的乘法。而 GEMM 應該是指 General Matrix Multiply,是矩陣對矩陣的乘法。

GEMV 在單筆推論時,速度比 GEMM 快,但是當 Context 較大量時速度比較慢。也就是說 Prefill 的 Token 量越大,越適合使用 GEMM 的矩陣乘法。因此若是要進行 LLM Service 部署,一般而言選擇 GEMM 即可。

註:根據此 Issue 所述,目前 vLLM 並不支援 GEMV 因為他只能在 Batch Size 1 的情況下使用。

量化工具

AutoAWQ 是個專門對 LLM 進行 AWQ 量化的工具,因為量化的過程是一層一層的量化,而且 AutoAWQ 可以將模型權重本身放在 CPU 裡面,只有需要量化時才把其中一層的權重放在 GPU Memory 裡面,因此記憶體的消耗非常少,甚至能用單張 RTX 3090 24GB 來量化 70B 規模的模型。

最近正值 PyTorch 更換 CUDA 版本的日子,所以套件之間的相依關係開始有些複雜,因此建議要用 AutoAWQ 的話,最好開一個新的 Conda 環境:

conda create -n AutoAWQ python=3.10
conda activate AutoAWQ
pip install autoawq==0.1.7  # 此為筆者最後測試的版本

安裝好 AutoAWQ 後,可以使用 AutoAWQForCausalLM 類別讀取原始模型:

import torch
from awq import AutoAWQForCausalLM

model_dir = "Models/Llama2-7B"

model = AutoAWQForCausalLM.from_pretrained(
    model_dir,
    torch_dtype=torch.float16,
    device_map="cpu",  # 權重可以放在 CPU 就好
    safetensors=True,  # 如果有支援 Safetensors 的話就開啟
)

這裡我們將模型放在 CPU RAM 裡面,避免 FP16 權重佔用太多記憶體,以至於 AutoAWQ 沒有空間進行量化運算。接下來,我們需要一份校準資料集,預設會使用 mit-han-lab/pile-val-backup 這份資料集進行校準,但我們也可以自行準備,只要把資料以 String List 的形式存放即可:

# 此為示範用的校準資料集,請根據自身應用替換成適當資料集
# 只要是 list[str] 的型態且 Token 數量要夠多就能量化
calib_data = ["Hello, World! " * 8] * 16

如果給的文本太短太少的話,進行量化時會發生 RuntimeError: torch.cat(): expected a non-empty list of Tensors 的錯誤。這是因為在 awq/utils/calib_data.py 第 30 行,作者放了一個 Token 數量必須大於 512 才會拿來用的條件,可能是因為原本 AWQ 的實做就是如此的緣故。

接下來放上相關設定,並開始進行量化:

from transformers import AutoTokenizer as TkCls

tokenizer = TkCls.from_pretrained(model_dir)

# 相關設定
quant_config = dict(
    zero_point=True,
    q_group_size=128,
    w_bit=4,
    version="GEMM",
)

# 進行量化
model.quantize(
    tokenizer,
    quant_config=quant_config,
    calib_data=calib_data,
)

# 儲存量化結果
output_dir = "Models/Llama2-7B-AWQ"
model.save_quantized(output_dir)
tokenizer.save_pretrained(output_dir)

筆者實測使用 RTX 3090 24GB 量化一個 7B 模型,大約只要七八分鐘就能完成。如果 CPU RAM 夠大,也能拿來量化 70B 的模型,大概要花費 100 分鐘左右。不過 70B 的模型即便量化到 4-Bit 也不能放在一張 24GB GPU 裡面使用就是了 QQ

這些量化完的模型,在 HF Transformers 裡面可以直接讀取進來使用:

from transformers import LlamaForCausalLM as ModelCls

model_dir = "Models/Llama2-7B-AWQ"
model: ModelCls = ModelCls.from_pretrained(
    model_dir,
    device_map="auto",
    use_safetensors=True,
)  # 約 4.5 GiB

不需要任何額外的設定就能讀取,雖然跟硬體效能有關,但 AWQ 模型在讀取速度上快滿多的。這種跟量化有關的技術,大神 TheBloke 同樣不會錯過。如果在 Hugging Face Hub 上搜尋 AWQ 的話,一字排開幾乎都是 TheBloke 上傳的模型。

從論文的描述來看,模型的精確度與校準資料集的內容相關性小的多,因此多數情況下直接拿別人上傳的 AWQ 模型來用理論上是沒問題的。但是實務上還是需要自行評測過才會知道是否沒問題,如果真的出問題了,或許就要考慮用自己的資料集來做校準量化了。

在推論的部份,與原本的用法並無二致:

from transformers import LlamaTokenizerFast as TkCls
from transformers import TextStreamer

tk: TkCls = TkCls.from_pretrained(model_dir)
ts = TextStreamer(tk)

prompt = "Hello, "
tokens = tk(prompt, return_tensors="pt").to("cuda")
model.generate(**tokens, max_new_tokens=16, streamer=ts)

實務上,更推薦使用 vLLMTGI 來運行 AWQ 模型,請參考各框架的相關文件。

vLLM AutoAWQ Usage | TGI CLI Quantize

參考


上一篇
LLM Note Day 32 - AutoGPTQ
系列文
LLM 學習筆記33
圖片
  直播研討會
圖片
{{ item.channelVendor }} {{ item.webinarstarted }} |
{{ formatDate(item.duration) }}
直播中

尚未有邦友留言

立即登入留言